# 从io中导入文件打开方法
from io import open
# 帮助使用正则表达式进行子目录的查询
import glob
import os
# 用于获得常见字母及字符规范化
import string
import unicodedata
# 导入随机工具random
import random
# 导入时间和数学工具包
import time
import math
# 导入torch工具
import torch
# 导入nn准备构建模型
import torch.nn as nn
# 引入制图工具包
import matplotlib.pyplot as plt
from torchsummary import summary

# 获取所有常用字符包括字母和常用标点
all_letters = string.ascii_letters + " .,;'"

# 获取常用字符数量
n_letters = len(all_letters)

# print("n_letter:", n_letters)
def unicodeToAscii(s):
    return ''.join(
        c for c in unicodedata.normalize('NFD', s)
        if unicodedata.category(c) != 'Mn'
        and c in all_letters
    )

# s = "Ślusàrski"
# a = unicodeToAscii(s)
# print(a)

data_path = "./data1/names/"

def readLines(filename):
    """从文件中读取每一行加载到内存中形成列表"""
    # 打开指定文件并读取所有内容, 使用strip()去除两侧空白符, 然后以'\n'进行切分
    lines = open(filename, encoding='utf-8').read().strip().split('\n')
    # 对应每一个lines列表中的名字进行Ascii转换, 使其规范化.最后返回一个名字列表
    return [unicodeToAscii(line) for line in lines]

# filename是数据集中某个具体的文件, 我们这里选择Chinese.txt
# filename = data_path + "Chinese.txt"
# lines = readLines(filename)
# print(lines)

# 构建的category_lines形如：{"English":["Lily", "Susan", "Kobe"], "Chinese":["Zhang San", "Xiao Ming"]}
category_lines = {}

# all_categories形如： ["English",...,"Chinese"]
all_categories = []

# 读取指定路径下的txt文件， 使用glob，path中可以使用正则表达式
for filename in glob.glob(data_path + '*.txt'):
    # 获取每个文件的文件名, 就是对应的名字类别
    category = os.path.splitext(os.path.basename(filename))[0]
    # 将其逐一装到all_categories列表中
    all_categories.append(category)
    # 然后读取每个文件的内容，形成名字列表
    lines = readLines(filename)
    # 按照对应的类别，将名字列表写入到category_lines字典中
    category_lines[category] = lines


# 查看类别总数
n_categories = len(all_categories)
print("n_categories:", n_categories)

# 随便查看其中的一些内容
# print(category_lines['Chinese'][:5])

def lineToTensor(line):
    """将人名转化为对应onehot张量表示, 参数line是输入的人名"""
    # 首先初始化一个0张量, 它的形状(len(line), 1, n_letters)
    # 代表人名中的每个字母用一个1 x n_letters的张量表示.
    tensor = torch.zeros(len(line), 1, n_letters)
    # 遍历这个人名中的每个字符索引和字符
    for li, letter in enumerate(line):
        # 使用字符串方法find找到每个字符在all_letters中的索引
        # 它也是我们生成onehot张量中1的索引位置
        tensor[li][0][all_letters.find(letter)] = 1
    # 返回结果
    return tensor

# line = "Bai"
# line_tensor = lineToTensor(line)
# print("line_tensor:", line_tensor)

# 构建RNN模型
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size,num_layers = 1):
        super(RNN,self).__init__()
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.input_size = input_size
        self.output_size = output_size

        # RNN的输入数据通常被组织为形状为 [batch_size, seq_len, input_size] ,隐藏层:形状为[num_layers, batch_size, hidden_size]
        self.rnn = nn.RNN(input_size, hidden_size, num_layers)
        # 实例化线性层,将rnn的输出维度转换为指定维度
        self.linear = nn.Linear(hidden_size, output_size)

        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, input1, hidden):
        # input1代表人名分类器中的输入张量, 形状为[1, n_letters] ,1代表batch_size
        # hidden代表隐藏层, 形状为[num_layers, batch_size, hidden_size]
        input1 = input1.unsqueeze(0)
        rr, hn = self.rnn(input1, hidden)  # 生成输出rr和新的隐藏层hn

        return self.softmax(self.linear(rr)) ,hn

    def initHidden(self):
        # 初始化隐藏层, 形状为[num_layers, batch_size, hidden_size]
        return torch.zeros(self.num_layers, 1, self.hidden_size)


class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size,num_layers = 1):
        super(LSTM, self).__init__()
        # 初始化参数
        self.hidden_size = hidden_size
        self.input_size = input_size
        self.output_size = output_size
        self.num_layers = num_layers
        # 实例化LSTM层
        self.lstm = nn.LSTM(input_size, hidden_size,num_layers)
        # 实例化线性层
        self.linear = nn.Linear(hidden_size, output_size)
        # 实例化softmax层,得到类别结果
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, input1, hidden, c):
        # input1代表人名分类器中的输入张量, 形状为[1, n_letters]
        # hidden代表隐藏层, 形状为[num_layers, batch_size, hidden_size]
        input1 = input1.unsqueeze(0)
        rr , (hn,cn) = self.lstm(input1, (hidden,c))  # 生成输出output和新的隐藏层hn
        return self.softmax(self.linear(rr)) ,hn,cn

    def initHiddenAndC(self):
        # 初始化隐藏层, 形状为[num_layers, batch_size, hidden_size]
        # hidden和c的形状为[num_layers, batch_size, hidden_size]
        c = hidden = torch.zeros(self.num_layers, 1, self.hidden_size)
        return hidden,c

class GRU(nn.Module):
    def __init__(self, input_size, hidden_size, output_size,num_layers = 1):
        super(GRU, self).__init__()
        # 初始化参数
        self.hidden_size = hidden_size
        self.input_size = input_size
        self.output_size = output_size
        self.num_layers = num_layers
        # 实例化GRU层
        self.gru = nn.GRU(input_size, hidden_size,num_layers)
        # 实例化线性层
        self.linear = nn.Linear(hidden_size, output_size)
        # 实例化softmax层,得到类别结果
        self.softmax = nn.LogSoftmax(dim=-1)

    def forward(self, input1, hidden):
        # input1代表人名分类器中的输入张量, 形状为[1, n_letters]
        # hidden代表隐藏层, 形状为[num_layers, batch_size, hidden_size]
        input1 = input1.unsqueeze(0)
        rr, hn = self.gru(input1, hidden)  # 生成输出output和新的隐藏层hn
        return self.softmax(self.linear(rr)), hn

    def initHidden(self):
        # 初始化隐藏层, 形状为[num_layers, batch_size, hidden_size]
        return torch.zeros(self.num_layers, 1, self.hidden_size)


input_size = n_letters     # 这个数量就直接决定了RNN模型中输入层的维度
# n_letters值就是RNN每个时间步输入向量的维度，每个字符都将被编码为一个在这个维度上的向量!
n_hidden = 128
# 输出尺寸为语言类别总数
output_size = n_categories

# input = lineToTensor("Bai")
# print(input.size())   -> torch.Size([3, 1, 57])
# input = lineToTensor("B")
# print(input.size())   -> torch.Size([1, 1, 57])
# 因为lineToTensor()函数返回的张量是三维的,所以需要将其转换为二维
# RNN的输入数据通常被组织为形状为 [batch_size, seq_len, input_size]


input1 = lineToTensor("B").squeeze(0)   # torch.Size([1, 57])
hidden = c = torch.zeros(1, 1, n_hidden)
rnn = RNN(input_size, n_hidden,output_size)  # n_letters表示输入的维度
lstm = LSTM(input_size, n_hidden,output_size)
gru = GRU(input_size, n_hidden,output_size)

rnn_output, next_hidden = rnn(input1, hidden)
print("rnn:", rnn_output)
print('rnn_shape', rnn_output.size())
print('-------------------------------------------------------')
lstm_output, next_hidden1, c = lstm(input1, hidden, c)
print("lstm:", lstm_output)
print('lstm_shape', lstm_output.size())
print('-------------------------------------------------------')
gru_output, next_hidden2 = gru(input1, hidden)
print("gru:", gru_output)
print('gru_shape', gru_output.size())




def categoryFromOutput(output):
    """从输出结果中获得指定类别, 参数为输出张量output"""
    # 从输出张量中返回最大的值和索引对象, 我们这里主要需要这个索引
    top_n, top_i = output.topk(1)
    # top_i对象中取出索引的值
    category_i = top_i[0].item()
    # 根据索引值获得对应语言类别, 返回语言类别和索引值
    return all_categories[category_i], category_i

category, category_i = categoryFromOutput(gru_output)
print("category:", category)
print("category_i:", category_i)


def randomTrainingExample():
    """该函数用于随机产生训练数据"""
    # 首先使用random的choice方法从all_categories随机选择一个类别
    category = random.choice(all_categories)
    # 然后在通过category_lines字典取category类别对应的名字列表
    # 之后再从列表中随机取一个名字
    line = random.choice(category_lines[category])
    # 接着将这个类别在所有类别列表中的索引封装成tensor, 得到类别张量category_tensor
    category_tensor = torch.tensor([all_categories.index(category)], dtype=torch.long)
    # 最后, 将随机取到的名字通过函数lineToTensor转化为onehot张量表示
    line_tensor = lineToTensor(line)
    return category, line, category_tensor, line_tensor

# # 我们随机取出十个进行结果查看
# for i in range(10):
#     category, line, category_tensor, line_tensor = randomTrainingExample()
#     print('category =', category, '/ line =', line, '/ category_tensor =', category_tensor)


# 定义损失函数为nn.NLLLoss，因为RNN的最后一层是nn.LogSoftmax, 两者的内部计算逻辑正好能够吻合
criterion = nn.NLLLoss()

learning_rate = 0.005
def trainRNN(category_tensor, line_tensor):    # category_tensor是二维张量,分别代表类别和索引
    """定义训练函数, 它的两个参数是category_tensor类别的张量表示, 相当于训练数据的标签,
       line_tensor名字的张量表示, 相当于对应训练数据"""

    # 在函数中, 首先通过实例化对象rnn初始化隐层张量
    hidden = rnn.initHidden()

    # 然后将模型结构中的梯度归0
    rnn.zero_grad()

    # 下面开始进行训练, 将训练数据line_tensor的每个字符逐个传入rnn之中, 得到最终结果
    for i in range(line_tensor.size()[0]):
        output, hidden = rnn(line_tensor[i], hidden)

    # 因为我们的rnn对象由nn.RNN实例化得到, 最终输出形状是三维张量, 为了满足于category_tensor
    # 进行对比计算损失, 需要减少第一个维度, 这里使用squeeze()方法
    loss = criterion(output.squeeze(0), category_tensor)
    # output.squeeze(0)两个维度, 第一个维度是1, 第二个维度为类别数量

    # 损失进行反向传播
    loss.backward()
    # 更新模型中所有的参数
    for p in rnn.parameters():
        # 将参数的张量表示与参数的梯度乘以学习率的结果相加以此来更新参数
        p.data.add_(p.grad.data, alpha=-learning_rate)
    # 返回结果和损失的值
    return output, loss.item()


# 与传统RNN相比多出细胞状态c

def trainLSTM(category_tensor, line_tensor):
    hidden, c = lstm.initHiddenAndC()
    lstm.zero_grad()
    for i in range(line_tensor.size()[0]):
        # 返回output, hidden以及细胞状态c
        output, hidden, c = lstm(line_tensor[i], hidden, c)
    loss = criterion(output.squeeze(0), category_tensor)
    loss.backward()

    for p in lstm.parameters():
        p.data.add_(p.grad.data,alpha=-learning_rate)
    return output, loss.item()


# 与RNN完全相同, 只不过名字改成了GRU

def trainGRU(category_tensor, line_tensor):
    hidden = gru.initHidden()
    gru.zero_grad()
    for i in range(line_tensor.size()[0]):
        output, hidden= gru(line_tensor[i], hidden)
    loss = criterion(output.squeeze(0), category_tensor)
    loss.backward()

    for p in gru.parameters():
        p.data.add_(p.grad.data,alpha=-learning_rate)
    return output, loss.item()




# 构建训练过程的日志打印函数
def timeSince(since):
    "获得每次打印的训练耗时, since是训练开始时间"
    # 获得当前时间
    now = time.time()
    # 获得时间差，就是训练耗时
    s = now - since
    # 将秒转化为分钟, 并取整
    m = math.floor(s / 60)
    # 计算剩下不够凑成1分钟的秒数
    s -= m * 60
    # 返回指定格式的耗时
    return '%dm %ds' % (m, s)


# 设置训练迭代次数
n_iters = 100000
# 设置结果的打印间隔
print_every = 50
# 设置绘制损失曲线上的制图间隔
plot_every = 10

def train(train_type_fn):
    """训练过程的日志打印函数, 参数train_type_fn代表选择哪种模型训练函数, 如trainRNN"""
    # 每个制图间隔损失保存列表
    all_losses = []
    # 获得训练开始时间戳
    start = time.time()
    # 设置初始间隔损失为0
    current_loss = 0
    # 从1开始进行训练迭代, 共n_iters次
    for iter in range(1, n_iters + 1):
        # 通过randomTrainingExample函数随机获取一组训练数据和对应的类别
        category, line, category_tensor, line_tensor = randomTrainingExample()
        # 将训练数据和对应类别的张量表示传入到train函数中
        output, loss = train_type_fn(category_tensor, line_tensor)
        # 计算制图间隔中的总损失
        current_loss += loss
        # 如果迭代数能够整除打印间隔
        if iter % print_every == 0:
            # 取该迭代步上的output通过categoryFromOutput函数获得对应的类别和类别索引
            guess, guess_i = categoryFromOutput(output)
            # 然后和真实的类别category做比较, 如果相同则打对号, 否则打叉号.
            correct = '✓' if guess == category else '✗ (%s)' % category
            # 打印迭代步, 迭代步百分比, 当前训练耗时, 损失, 该步预测的名字, 以及是否正确
            print('%d %d%% (%s) %.4f %s / %s %s' % (iter, iter / n_iters * 100, timeSince(start), loss, line, guess, correct))

        # 如果迭代数能够整除制图间隔
        if iter % plot_every == 0:
            # 将保存该间隔中的平均损失到all_losses列表中
            all_losses.append(current_loss / plot_every)
            # 间隔损失重置为0
            current_loss = 0
    # 返回对应的总损失列表和训练耗时
    return all_losses, int(time.time() - start)



# # 开始训练RNN, LSTM, GRU模型并制作对比图
# # 调用train函数, 分别进行RNN, LSTM, GRU模型的训练
# # 并返回各自的全部损失, 以及训练耗时用于制图
# all_losses1, period1 = train(trainRNN)
# all_losses2, period2 = train(trainLSTM)
# all_losses3, period3 = train(trainGRU)
#
# # 绘制损失对比曲线, 训练耗时对比柱张图
# # 绘制损失对比曲线
# plt.figure(figsize=(12, 6))  # 可选：设定图形大小
# plt.plot(all_losses1, label="RNN", linewidth=2)
# plt.plot(all_losses2, color="red", label="LSTM", linewidth=2)
# plt.plot(all_losses3, color="orange", label="GRU", linewidth=2)
# plt.title("Loss Comparison Over Training Epochs")
# plt.xlabel("Epochs")
# plt.ylabel("Loss")
# plt.legend(loc='upper right')
# plt.show()  # 显示图像并暂停程序执行，直到图像关闭
#
# # 绘制训练耗时对比柱状图
# plt.figure(figsize=(8, 6))
# x_data = ["RNN", "LSTM", "GRU"]
# y_data = [period1, period2, period3]
# plt.bar(x_data, y_data, color=['blue', 'red', 'orange'])
# plt.title("Training Duration Comparison")
# plt.ylabel("Duration (seconds/minutes/hours)")  # 根据period单位调整
# plt.show()



def evaluateRNN(line_tensor):
    """评估函数, 和训练函数逻辑相同, 参数是line_tensor代表名字的张量表示"""
    # 初始化隐层张量
    hidden = rnn.initHidden()
    # 将评估数据line_tensor的每个字符逐个传入rnn之中
    for i in range(line_tensor.size()[0]):
        output, hidden = rnn(line_tensor[i], hidden)
    # 获得输出结果
    return output.squeeze(0)


def evaluateLSTM(line_tensor):
    # 初始化隐层张量和细胞状态张量
    hidden, c = lstm.initHiddenAndC()
    # 将评估数据line_tensor的每个字符逐个传入lstm之中
    for i in range(line_tensor.size()[0]):
        output, hidden, c = lstm(line_tensor[i], hidden, c)
    return output.squeeze(0)



def evaluateGRU(line_tensor):
    hidden = gru.initHidden()
    # 将评估数据line_tensor的每个字符逐个传入gru之中
    for i in range(line_tensor.size()[0]):
        output, hidden = gru(line_tensor[i], hidden)
    return output.squeeze(0)


# line = 'Bai'
# line_tensor = lineToTensor(line)
# rnn_output = evaluateRNN(line_tensor)
# lstm_output = evaluateLSTM(line_tensor)
# gru_output = evaluateGRU(line_tensor)
# print("RNN:", rnn_output)
# print("LSTM:", lstm_output)
# print("GRU:", gru_output)


# 构建预测函数
def predict(input_line, evaluate, n_predictions=3):
    """预测函数, 输入参数input_line代表输入的名字,
       n_predictions代表需要取最有可能的top个"""
    # 首先打印输入
    print('\n> %s' % input_line)

    # 以下操作的相关张量不进行求梯度
    with torch.no_grad():
        # 使输入的名字转换为张量表示, 并使用evaluate函数获得预测输出
        output = evaluate(lineToTensor(input_line))

        # 从预测的输出中取前3个最大的值及其索引
        topv, topi = output.topk(n_predictions, 1, True)
        # 创建盛装结果的列表
        predictions = []
        # 遍历n_predictions
        for i in range(n_predictions):
            # 从topv中取出的output值
            value = topv[0][i].item()
            # 取出索引并找到对应的类别
            category_index = topi[0][i].item()
            # 打印ouput的值, 和对应的类别
            print('(%.2f) %s' % (value, all_categories[category_index]))
            # 将结果装进predictions中
            predictions.append([value, all_categories[category_index]])

for evaluate_fn in [evaluateRNN, evaluateLSTM, evaluateGRU]:
    print("-"*18)
    predict('Dovesky', evaluate_fn)
    predict('Jackson', evaluate_fn)
    predict('Satoshi', evaluate_fn)
    print("-"*18)



